import os
import numpy as np
import cv2


def get_color_map(N=256):
    def bitget(byteval, idx):
        return ((byteval & (1 << idx)) != 0)

    cmap = np.zeros((N, 3), dtype=np.uint8)
    for i in range(N):
        r = g = b = 0
        c = i
        for j in range(8):
            r = r | (bitget(c, 0) << 7-j)
            g = g | (bitget(c, 1) << 7-j)
            b = b | (bitget(c, 2) << 7-j)
            c = c >> 3

        cmap[i] = np.array([r, g, b])

    return cmap


def shift_mask(mask, ZY):
    if ZY == "ZY20210800001":
        dx, dy = 0, 0
    elif (ZY == "ZY20210800002") or (ZY == "ZY20210800004"):
        dx, dy = 15, -3
    elif ZY == "ZY20210800003":
        dx, dy = 15, -30
    else:
        raise NotImplementedError

    rows, cols, _ = mask.shape
    MAT = np.float32([[1, 0, dx], [0, 1, dy]])
    dst = cv2.warpAffine(mask, MAT, (cols, rows))
    return dst


def read_rgb_and_gt_mask(video_dir, idx):
    """
    input: video_dir(e.g. "/share/datasets/HOI4D_overall/ZY20210800003/H3/C12/N30/S176/s02/T1"), idx(0-299)
    return: rgb, mask
    """

    # extract ZY
    p1 = video_dir.find("ZY")
    p2 = video_dir[p1:].find("/")
    ZY = video_dir[p1:p1+p2]
    if ZY == "ZY20210800001":
        rgb_path = os.path.join(video_dir, "align_image", str(idx) + ".jpg")
    else:
        rgb_path = os.path.join(video_dir, "shift_rgb", str(idx) + ".jpg")
    
    # read data and shift
    mask_path = os.path.join(video_dir, "refine_2Dseg", "mask", str(idx).zfill(5) + ".png")  # unshifted
    rgb = cv2.imread(rgb_path, cv2.IMREAD_UNCHANGED)[:, :, ::-1]
    mask = cv2.imread(mask_path, cv2.IMREAD_UNCHANGED)[:, :, ::-1]  # unshifted
    mask = shift_mask(mask, ZY)  # shifted

    return rgb, mask


if __name__ == "__main__":
    video_dir = "/share/datasets/HOI4D_overall/ZY20210800003/H3/C12/N30/S176/s02/T1"
    idx = 100
    rgb, mask = read_rgb_and_gt_mask(video_dir, idx)
